﻿# LLM的输出格式成Python list形式，类似['a', 'b', 'c']
from dotenv import load_dotenv
from langchain.llms import OpenAI
from langchain.output_parsers import PydanticOutputParser, CommaSeparatedListOutputParser
from langchain.prompts import PromptTemplate
from langchain.pydantic_v1 import BaseModel, Field, validator
from typing import List
import os

load_dotenv("../ai.env")

api_base = os.getenv("OPENAI_API_BASE")
api_key = os.getenv("OPENAI_KEY")

# 构造一个llm
model = OpenAI(
    model="gpt-3.5-turbo-instruct",
    temperature=0,
    openai_api_key=api_key,
    open_api_base=api_base
)

parser = CommaSeparatedListOutputParser()

prompt = PromptTemplate(
    template="列出5个{subject}.\n{format_instructions}",
    input_variables=["subject"],
    partial_variables={"format_instructions": parser.get_format_instructions()}
)

_input = prompt.format(subject="常见的小狗的名字")
output = model(_input)
print(output)
# 格式化
print(parser.parse(output))
